import os
import base64
from array import array
from collections import namedtuple
from collections import defaultdict
import xml.etree.ElementTree as et
from xml.parsers.expat import ExpatError
from easier_raw_class import RawFileExtractor, MockExtractor
#
#
#
def decode(text):
    "decodifica el texto extraido de una file mzXML que codifica el espectro "
    spec = array('f')
    spec.fromstring(base64.b64decode(text))
    spec.byteswap()
    return spec
#
#
def format_espectro(espectro):
    """
    lee un espectro en formato lineal [mz1, int1, mz2, int2, ...]
    y lo devuelve  como ['mz1 int1','mz2 int2','mz3 int3',...]
    """
    splist = []
    i = 0
    while True:
        try:
            intsty = espectro[i+1]
        except IndexError:
            break
        
        splist.append('%s %s' %(espectro[i], intsty))
        i += 2
        
    return splist
#
#
def iter_mzxml_data(xml, add_quant=True):
    """iterates over scan elements in ReAdW mzXML file.
    
    yield spectrum and params ('scan, mslevel, rt, charge, method, mz, spec')

    the process assumes a scan series of the type:
     
        fscan + (qscan * n + (cidscan + (cidms3 * y))) * m
            m  numero de precursores seleccionados (m>0)
            n  average de espectros de cuantificacion utilizados (n>=0)
            y  numero de ms3 (y >= 0)
    """
    #
    try:
        doc = et.parse(xml)
    except ExpatError, e:
        print "archivo %s produce error\n'%s'" % (xml, e)
        return
    #
    root = doc.getroot()
    children = root.getchildren()
    msRun = children[0]
    #start_time = msRun.attrib.get('startTime')
    #start_time = float(start_time[2:-1])
    #print start_time
    for scanelement in msRun.getchildren():
        if scanelement.tag.endswith('scan'):
            mgfdatalist = process(scanelement)
        else:
            continue
                    
        quantlist = []

        for scan in mgfdatalist:
           
            spec_data = decode(scan.spec)
            spec = format_espectro(spec_data)
            scan_number =  scan.scan
            
            if (scan.method == 'HCD') or (scan.method == 'PQD'):
                #el primer scan de cuantificacion
                if not quantlist:
                    start_qscan = scan_number
                    quantlist = [spec]
                # los siguientes qscans
                else:
                    quantlist.append(spec)
                    
            elif scan.method == 'CID':
                pepmass = scan.mz
                rt = float(scan.rt[2:-1])
                rt = '%.4f' % (rt/60)
                
                #paso a cid despues de recoger varios qscans
                if quantlist:
                    qspec = average(quantlist)
                    
                    scan_range = "%s-%s" %(start_qscan, scan_number)
                    spec = add_quant_to_cid(qspec, spec, 0)
                    quant = []
                # no hay qscans o es ms3
                # ojo, ms3 NO debe sumar tambien quant, eso lo hace integrator?
                # no seria mejor hacerlo aqui?
                # no seria mejor asignar una quantificacion a cada numero de espectro
                # llevarla separada y ligarla al final, despues de la identificacion?,
                # no afectaria la busqueda.
                else:
                    scan_range = scan_number
                    
                   
                params = (0, scan.charge, rt,
                          scan.mslevel, scan_range, scan.mz )
                                
                yield spec, params
                
                
def average(speclist):
    """extrae los datos de quantificacion de una lista de espectros
    
    no hace caso de si es itraq o tmt, solo suma iones. la cuantificacion se
    hace en otra parte. en principio en el integrator aunque no es una buena idea
    porque integrator lo hace a saco sin tener en cuenta la precision de masa y
    otras finuras (que estaban resueltas en psearch!, mierda!!)
    
    los espectros son de la forma
    [['mz1 int1','mz2 int2','mz3 int3',...], ....] 
    
    """
   
    if len(speclist) > 1:
        specdict = defaultdict(float)
        
        for spec in speclist:
            for value in spec:
                mass, int = value.split()
                specdict[float(mass)] += float(int)
                
        
        masses = sorted(specdict.keys())
        
        spec = ['%f %f' %(mass, specdict[mass]) for mass in masses]
        
    else:
        spec = speclist[0]
    
    return spec
#
#
def build_mgfs_from_mzxml(xml, mgf, temp, remove_z_one, split_2_3, add_quant):
    """
    build custom mgf from ReAdW mzXML
    """
    xmlpath = os.path.join(temp, xml)
    (xmlbasename, _) = xml.rsplit('.')
    
    fhndls = get_file_handles(temp, mgf, split_2_3, remove_z_one)
    
    if split_2_3 and remove_z_one:
        (fh12, fh13, fh2, fh3) = fhndls
    elif split_2_3:
        (fh2, fh3) = fhndls
    elif remove_z_one:
        (fh123, fh231) = fhndls
    else:
        (fh,) = fhndls
    
    for espectro, (mass, charge, rt, ms, sc, mz) in iter_mzxml_data(xmlpath):
        head_hcd = ['']
        head_hcd.extend(espectro)
        head_hcd.append('END IONS\n\n')
        espectro = build_mgf(head_hcd, mass, charge, xmlbasename, rt, ms, sc, mz)
        espectro[0] = espectro[0].rstrip()
        espectro = '\n'.join(espectro)
        if split_2_3 and remove_z_one:
            if charge == '1':
                if ms == '2':
                    fh12.write(espectro)
                elif ms == '3':
                    fh13.write(espectro)
            else:
                if ms == '2':
                    fh2.write(espectro)
                elif ms == '3':
                    fh3.write(espectro)
        elif split_2_3:
            if ms == '2':
                fh2.write(espectro)
            elif ms == '3':
                fh3.write(espectro) 
        elif remove_z_one:
            if charge == '1':
                fh123.write(espectro)
            else:
                fh231.write(espectro)
        else:
            fh.write(espectro)
            
    for hndl in fhndls: hndl.close()
#
#
def process(scanelement):
    """
    devuelve una lista tipo:
    [pqd, pqd, cid, ms3, ..., pqd, pqd, cid, ms3]
    """
   
    if scanelement.attrib.get('msLevel') == '1':
        childs = scanelement.getchildren()
        msms = [item for item in childs if item.tag.endswith('scan')]
    else:
        print 'return. es msLevel ', scanelement.attrib.get('msLevel')
        return
    
    msmsdata = []
    for ms in msms:
        data = get_scandata(ms)
        msmsdata.append(data)
        ms3 = [ms for ms in ms.getchildren() if ms.tag.endswith('scan')]
        for ms in ms3:
            data = get_scandata(ms)
            msmsdata.append(data)
            
    return msmsdata
#
#
def get_scandata(ms):
    "return named tuple of scan data"
    ScanData = namedtuple('ScanData',
                          'scan, mslevel, rt, charge, method, mz, spec')
    
    mschilds = ms.getchildren()
    return ScanData(ms.attrib.get('num'),
                    ms.attrib.get('msLevel'),
                    ms.attrib.get('retentionTime'), 
                    mschilds[0].attrib.get('precursorCharge'),
                    mschilds[0].attrib.get('activationMethod'),
                    mschilds[0].text,
                    mschilds[1].text)
#
#
def build_mgf_from_dtas(dtadir, raw, outfile, remove_z_one, split_2_3, add_quant):
    """A partir de dta que contiene unicamente:
    
    1) filename
    (filename.scaninicial.scanfinal.carga)
    'test_DIGE.120.125.2.dta'
    'test_DIGE.105.105.2.dta'     (no agrupado)          
    
    2) cabecera dta
    (pept_Mr carga)
    996.70949 2                             
    
    genera mgf con cabecera tipo:
    
    BEGIN IONS
    TITLE=test_DIGE,Scan:120-125,MS:2,Rt:1833.494
    PEPMASS=419.32                          (pept m/z ?)
    SCANS=120-125
    RTINSECONDS=1833.494
    
    Used by ExtracMsn classes in easymgf
    """
    
    dtas = os.listdir(dtadir)
    dtas = [arch for arch in dtas if arch.endswith('.dta')] 
    dtas.sort(key=key_sort)
    
    rawdata = RawFileExtractor(raw)
    
    tempdir, outfile = os.path.split(outfile)
    
    if split_2_3:
        outhandle_2 = fhopen(tempdir, 'ms2_', outfile)
        outhandle_3 = fhopen(tempdir, 'ms3_', outfile)
    else:
        outhandle = fhopen(tempdir, '', outfile)
        
    
    hay_hcd = False
    
   
    
    for dta in dtas:
        
        try:
            (rawname, scan_ini, scan_end, charge, ext) = dta.split('.')
        except ValueError:
            print 'error'
            continue
        
        #print rawname
        
        if remove_z_one and (charge == 1): continue
        
        if scan_ini == scan_end:
            group_scans = scan_ini
        else:
            group_scans = '%s-%s' % (scan_ini, scan_end)
            
        #continue if scan has been already used  ??
        #FTMS + c NSI d Full ms2 1070.52@hcd59.00 [110.00-2000.00]
            
        dtapath = os.path.join(dtadir, dta)
        
        handle = open(dtapath)
        
        lines = handle.readlines()
        cabecera = lines[0]
        (mass, charge) = cabecera.strip().split() 
        (mode, ms) = rawdata.get_mode_and_ms_level(scan_end)
        rt = rawdata.get_rt(scan_end)
        ms = ms[-1]
                
        if (mode == 'hcd') and (add_quant == True):
            #print 'hcd en ', group_scans
            hcd_mass_charge = (mass, charge)
            hcd_dta = lines[:]
            hay_hcd = True
            continue

        if hay_hcd and ((mass, charge) == hcd_mass_charge):
            lines = add_quant_to_cid(hcd_dta, lines, 1)
            lines = build_mgf(lines, mass, charge, rawname, rt, ms, group_scans)
            hay_hcd = False
        elif hay_hcd:
            lines = build_mgf(hcd_dta, mass, charge, rawname, rt, ms, group_scans)
            
            if split_2_3:
                if ms == '2':
                    outhandle_2.write(''.join(lines) + '\n')
                elif ms == '3':
                    outhandle_3.write(''.join(lines) + '\n')
            else:        
                outhandle.write(''.join(lines)+'\n')
            
            lines = build_mgf(lines, mass, charge, rawname, rt, ms, group_scans)
            hay_hcd = False
        else:
            lines = build_mgf(lines, mass, charge, rawname, rt, ms, group_scans)
                
        if split_2_3:
            if ms == '2':
                outhandle_2.write(''.join(lines) + '\n')
            elif ms == '3':
                outhandle_3.write(''.join(lines) + '\n')
        else:        
            outhandle.write(''.join(lines)+'\n')
        
    
    rawdata.close()
    
    if split_2_3:
        outhandle_2.close()
        outhandle_3.close()
    else:
        outhandle.close()
#
#
def add_quant_to_cid(hcd, cid, startline):
    "Mix HCD/PQD quantitative ions with CID scan."
    for indx, hcd_line in enumerate(hcd[startline:]):
        mass , resto = hcd_line.split()
        if float(mass) > 150:
            last_line = indx + startline
            break
        
    for indx, cid_line in enumerate(cid[startline:]):
        mass , resto = cid_line.split()
        if float(mass) > 150:
            first_line = indx + startline
            break    
    
    try:
        newdta = cid[:startline] + hcd[startline:last_line] + cid[first_line:]
    #sometimes the spectrum have no signals and
    #lastline and/or first_line are not assigned
    except UnboundLocalError:
        newdta = ['0 0']
        
    return newdta
#
#
def build_mgf(lines, mass, charge, rawname, rt, ms, group_scans, mz=None):
    """Build and add a custom mgf header to an spectrum
    
    the spectrum is in the form of a list of lines with first line available
    for the new header (no info must be lost)    
    """
    FORMAT = ('BEGIN IONS\n'
              'TITLE=%s,Scan:%s,MS:%s,Rt:%s\n'
              'PEPMASS=%s\n'
              'SCANS=%s\n'
              'RTINSECONDS=%s\n')
                    
    if not mz:
        mz = (float(mass) + (int(charge) - 1) * 1.007825) / int(charge)
    
    text = FORMAT % (rawname, group_scans, ms, rt,
                      mz,
                      group_scans,
                      rt)
                    
    lines[0] = text
    
    return lines
#
#
def key_sort(name):
    "Key for sorting dtas on the basis on the scan number in the name"
    splited = name.split('.')
    return int(splited[1])
#
#
def fhopen(fdir, prefix, fname):
    "for safe openning of new files in append mode"
    fpath = os.path.join(fdir, prefix + fname)
    if os.path.exists(fpath):
        os.remove(fpath)
        
    return open(fpath, 'a')
#
#
def get_file_handles(temp, fname, split_2_3, remove_z_one):
    "Return the filehandles to the different output files"
    if split_2_3 and remove_z_one:
        fh12 = fhopen(temp, '1ms2_', fname)
        fh13 = fhopen(temp, '1ms3_', fname)
        fh2  = fhopen(temp, '2ms21_', fname)
        fh3  = fhopen(temp, '2ms31_', fname)
        return (fh12, fh13, fh2, fh3)
    
    elif split_2_3:
        fh2 = fhopen(temp, 'ms2_', fname)
        fh3 = fhopen(temp, 'ms3_', fname)
        return (fh2, fh3)
    
    elif remove_z_one:
        fh123 = fhopen(temp, '1ms23_', fname)
        fh231 = fhopen(temp, 'ms231_', fname)
        return (fh123, fh231)
    
    else:
        fh = fhopen(temp, 'out_', fname)
        return (fh,)
#
#
# 
if __name__ == "__main__":
        
    dtadir = 'C:/Python26/programas/easier/extract_msn_temp/test_DIGE'
    raw = 'C:/Python26/programas/easier/test_mgf/test_DIGE.RAW'
    outfile = 'C:/Python26/programas/easier/test_mgf/mgf_from_dta'
    
    build_mgf_from_dtas(dtadir, raw, outfile)
        
    